from torch.utils.data import Dataset
import pandas as pd
import numpy as np
import torch


class TurbineDataset(Dataset):
    '''
    base：表示时间序列预测中的“上一个点”
    trend：趋势项
    season：周期项
    y：目标值
    seq：历史窗口长度
    '''
    def __init__(self, base:pd.DataFrame, trend:pd.DataFrame, season:pd.DataFrame, y:pd.Series, seq:int, device):
        super(TurbineDataset, self).__init__()
        self.base = torch.tensor(base.to_numpy(dtype=np.float32), dtype=torch.float32).to(device)
        self.trend = torch.tensor(trend.to_numpy(dtype=np.float32), dtype=torch.float32).to(device)
        self.season = torch.tensor(season.to_numpy(dtype=np.float32), dtype=torch.float32).to(device)
        self.y = torch.tensor(y.to_numpy(dtype=np.float32), dtype=torch.float32).to(device)
        self.seq = seq
        self.device = device


    def __getitem__(self, index):
        return self.base[index: index + self.seq], self.trend[index: index + self.seq], self.season[index: index + self.seq], self.y[index + self.seq]


    def __len__(self):
        return len(self.y) - self.seq